package adapters

import (
	"ai-gateway/common"
	"ai-gateway/db"
	"ai-gateway/proto"
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
)

// DeepSeekStreamAdapter 实现流式接口，处理DeepSeek的SSE响应
type DeepSeekStreamAdapter struct{}

// StreamCall流式调用DeepSeek API，通过gRPC流返回实时结果
func (d *DeepSeekStreamAdapter) StreamCall(ctx context.Context, req *proto.AIRequest,
	stream proto.AdapterService_StreamCallAIServer) error {
	// 从数据库获取DeepSeek配置
	var config struct {
		ApiKey     string  `db:"api_key"`
		ApiUrl     string  `db:"api_url"`
		TokenPrice float64 `db:"token_price"`
	}
	err := db.DB.Get(&config, "SELECT api_key, api_url, token_price FROM provider_configs WHERE provider = 'deepseek'")
	if err != nil {
		return fmt.Errorf("未找到DeepSeek配置: %v", err)
	}

	// 构造流式请求体
	requestBody := map[string]interface{}{
		"model": "deepseek-chat", // 正确模型名
		"messages": []map[string]string{ // 与同步调用一致，使用messages数组
			{"role": "user", "content": req.Prompt},
		},
		"stream":     true, // 启用流式响应
		"max_tokens": 1024,
	}
	jsonBody, err := json.Marshal(requestBody)
	if err != nil {
		return fmt.Errorf("请求体序列化失败: %v", err)
	}

	// 发起HTTP请求
	httpReq, err := http.NewRequestWithContext(ctx, "POST", config.ApiUrl, strings.NewReader(string(jsonBody)))
	if err != nil {
		return fmt.Errorf("创建请求失败: %v", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+config.ApiKey)

	client := &http.Client{}
	resp, err := client.Do(httpReq)
	if err != nil {
		return fmt.Errorf("调用API失败: %v", err)
	}
	defer resp.Body.Close()

	// 检查HTTP状态码
	if resp.StatusCode != http.StatusOK {
		body, _ := json.Marshal(resp.Body)
		return fmt.Errorf("API返回错误（状态码: %d）: %s", resp.StatusCode, string(body))
	}

	// 解析SSE流式响应（使用bufio.Scanner逐行处理，更可靠）
	scanner := bufio.NewScanner(resp.Body)
	var (
		totalTokens int
		fullResult  strings.Builder
	)

	for scanner.Scan() {
		line := scanner.Text()
		if line == "" {
			continue
		}

		// 处理SSE格式（"data: {json}\n" 或 "data: [DONE]"）
		if !strings.HasPrefix(line, "data: ") {
			continue // 忽略非数据行
		}
		data := strings.TrimPrefix(line, "data: ")

		// 检查流结束标记
		if data == "[DONE]" {
			break
		}

		// 解析流式响应片段
		var streamResp struct {
			Choices []struct {
				Delta struct {
					Content string `json:"content"` // 实时返回的文本片段
				} `json:"delta"`
				FinishReason *string `json:"finish_reason"` // 可能为null，需用指针
			} `json:"choices"`
			Usage struct {
				TotalTokens int `json:"total_tokens"` // 累计Token数（部分接口在结束时返回）
			} `json:"usage"`
		}

		if err := json.Unmarshal([]byte(data), &streamResp); err != nil {
			common.Logger.Warn().Err(err).Str("data", data).Msg("解析流式片段失败")
			continue
		}

		// 提取文本片段并累加
		var segment string
		for _, choice := range streamResp.Choices {
			segment += choice.Delta.Content
			// 结束时更新总Token数（部分接口在最后一次响应返回usage）
			if choice.FinishReason != nil && *choice.FinishReason == "stop" {
				totalTokens = streamResp.Usage.TotalTokens
			}
		}

		// 通过gRPC流返回实时片段
		if segment != "" {
			fullResult.WriteString(segment)
			if err := stream.Send(&proto.AIResponse{
				Result:     segment,
				Provider:   "deepseek",
				TokenUsed:  int32(totalTokens),
				Cost:       int64(float64(totalTokens) * config.TokenPrice * 100),
				IsFinished: false,
			}); err != nil {
				return fmt.Errorf("发送流式响应失败: %v", err)
			}
		}
	}

	// 检查扫描错误
	if err := scanner.Err(); err != nil {
		return fmt.Errorf("读取流式响应失败: %v", err)
	}

	// 发送最终完整结果（标记结束）
	return stream.Send(&proto.AIResponse{
		Result:     fullResult.String(),
		Provider:   "deepseek",
		TokenUsed:  int32(totalTokens),
		Cost:       int64(float64(totalTokens) * config.TokenPrice * 100),
		IsFinished: true,
	})
}

// 注册流式适配器
func init() {
	RegisterStreamAdapter("deepseek", &DeepSeekStreamAdapter{})
}
